# frequency_analysis.py
# Harmonic frequency analysis to confirm true minimum
# Verifies optimized geometry has no imaginary frequencies

from pyscf import gto, dft
from pyscf.hessian import rks
from pyscf.solvent import pcm
import numpy as np

print("="*70)
print("HARMONIC FREQUENCY ANALYSIS")
print("="*70)
print()

# Load optimized geometry from XYZ file
print("Loading optimized geometry from camphorquinone_optimized.xyz...")
with open('camphorquinone_optimized.xyz', 'r') as f:
    lines = f.readlines()

# Parse XYZ format: skip first two lines (number of atoms + comment)
xyz_lines = lines[2:]
atom_list = []
for line in xyz_lines:
    parts = line.split()
    if len(parts) >= 4:
        atom = parts[0]
        x, y, z = map(float, parts[1:4])
        atom_list.append([atom, (x, y, z)])

print(f"Loaded {len(atom_list)} atoms")
print()

# Build molecule
mol = gto.Mole()
mol.atom = atom_list
mol.basis = '6-31g(d)'
mol.charge = 0
mol.spin = 0
mol.verbose = 4
mol.build()

# DFT with PCM solvent
print("Setting up B3LYP/6-31G(d)/PCM(ε=4.0) calculation...")
mf = dft.RKS(mol)
mf.xc = 'B3LYP'
mf = pcm.PCM(mf)
mf.with_solvent.eps = 4.0
mf.verbose = 4

print()
print("Running SCF calculation...")
E_scf = mf.kernel()
print(f"SCF Energy: {E_scf:.10f} Hartree")
print()

# Compute Hessian
print("="*70)
print("COMPUTING HESSIAN MATRIX")
print("(This may take several minutes...)")
print("="*70)
print()

hess_obj = rks.Hessian(mf)
hess = hess_obj.kernel()

# Save Hessian matrix
np.save('hessian.npy', hess)
print("✓ Hessian matrix saved to hessian.npy")
print()

# Compute frequencies from Hessian
print("="*70)
print("COMPUTING VIBRATIONAL FREQUENCIES")
print("="*70)
print()

# Mass-weight the Hessian and diagonalize
mass = mol.atom_mass_list()
natm = mol.natm

# Create mass-weighted Hessian
hess_mw = np.zeros_like(hess)
for i in range(natm):
    for j in range(natm):
        hess_mw[3*i:3*i+3, 3*j:3*j+3] = hess[3*i:3*i+3, 3*j:3*j+3] / np.sqrt(mass[i] * mass[j])

# Diagonalize to get frequencies
eigenvalues, eigenvectors = np.linalg.eigh(hess_mw)

# Convert eigenvalues to frequencies (cm^-1)
# Factor = sqrt(Eh/(amu * bohr^2)) * (c in cm/s) / (2*pi)
# = 5140.48 cm^-1 per sqrt(Eh/amu)
au_to_cm = 5140.48756
frequencies = np.sign(eigenvalues) * np.sqrt(np.abs(eigenvalues)) * au_to_cm

# Count imaginary frequencies
n_imaginary = np.sum(frequencies < 0)
real_freqs = frequencies[frequencies > 10.0]  # Exclude near-zero modes

print(f"Total modes: {len(frequencies)}")
print(f"Imaginary frequencies: {n_imaginary}")
print(f"Real frequencies (>10 cm⁻¹): {len(real_freqs)}")
print()

if n_imaginary == 0:
    print("✓ NO IMAGINARY FREQUENCIES FOUND")
    print("  → Geometry is a TRUE LOCAL MINIMUM")
    print()
else:
    print(f"⚠ WARNING: {n_imaginary} imaginary frequencies detected")
    print("  → Geometry may be a saddle point or transition state")
    print()

# Display frequency summary
print("="*70)
print("FREQUENCY SUMMARY")
print("="*70)
print()
print("Lowest 10 real frequencies (cm⁻¹):")
sorted_real = np.sort(real_freqs)
for i, freq in enumerate(sorted_real[:10], 1):
    print(f"  {i:2d}. {freq:8.2f} cm⁻¹")

print()
print("Highest 5 frequencies (cm⁻¹):")
for i, freq in enumerate(sorted_real[-5:], len(sorted_real)-4):
    print(f"  {i:2d}. {freq:8.2f} cm⁻¹")

# Identify C=O stretches (typically 1700-1750 cm^-1)
co_stretches = real_freqs[(real_freqs > 1650) & (real_freqs < 1800)]
print()
print(f"C=O stretching modes ({len(co_stretches)} found):")
for freq in sorted(co_stretches):
    print(f"  {freq:.2f} cm⁻¹")

print()
print("="*70)
print("FREQUENCY ANALYSIS COMPLETE")
print("="*70)
print()
print("Output files:")
print("  ✓ hessian.npy")
print()
print("Summary:")
if n_imaginary == 0:
    print("  ✓ Geometry validated as true minimum")
    print("  ✓ Ready for TD-DFT calculations")
else:
    print("  ✗ Geometry needs further optimization")
print()
